import cv2
import pickle
import numpy as np
import matplotlib.pyplot as plt

def load_model(model_path):
    """加载保存的模型"""
    with open(model_path, "rb") as file:
        clf = pickle.load(file)
    return clf

def load_images(original_image_path, segment_image_path):
    """加载原图和分割图"""
    original_image = cv2.imread(original_image_path, 0)
    segment_image = cv2.imread(segment_image_path, 0)
    if original_image is None or segment_image is None:
        raise FileNotFoundError("原图或分割图未找到")
    return original_image, segment_image

def add_features(image):
    """添加额外的特征 - 均值滤波"""
    mean_filtered = cv2.blur(image, (5, 5))
    features = np.concatenate((image.reshape(-1, 1), mean_filtered.reshape(-1, 1)), axis=1)
    return features

def predict_segmentation(clf, features, image_shape):
    """对新的特征矩阵进行预测"""
    predicted_segment = clf.predict(features[:, [0]])  # 仅使用原始图像特征进行预测
    predicted_segment = predicted_segment.reshape(image_shape)
    return predicted_segment

def calculate_accuracy(predicted_segment, true_segment):
    """计算准确率"""
    accuracy = (predicted_segment == true_segment).mean()
    return accuracy

def display_images(original_image, segment_image, predicted_segment):
    """显示原图、原分割图和推理分割图"""
    fig, axes = plt.subplots(1, 3, figsize=(12, 4))

    axes[0].imshow(original_image, cmap='gray')
    axes[0].set_title("Original Image")

    axes[1].imshow(segment_image, cmap='gray')
    axes[1].set_title("Segmented Image")

    axes[2].imshow(predicted_segment, cmap='gray')
    axes[2].set_title("Predicted Segmentation")

    plt.show()

def main():
    model_path = r"D:/tuxiang/stone/clf.pkl"
    original_image_path = r"D:/tuxiang/stone/Sandstone_2.tif"
    segment_image_path = r"D:/tuxiang/stone/Sandstone_2_segment.tif"

    # 加载模型
    clf = load_model(model_path)

    # 加载图像
    original_image, segment_image = load_images(original_image_path, segment_image_path)

    # 添加特征
    features = add_features(original_image)

    # 预测分割
    predicted_segment = predict_segmentation(clf, features, original_image.shape)

    # 计算准确率
    accuracy = calculate_accuracy(predicted_segment, segment_image)
    print("准确率:", accuracy)

    # 显示图像
    display_images(original_image, segment_image, predicted_segment)

if __name__ == "__main__":
    main()